package org.finesys.common.dataPermission.db;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.Collection;
import java.util.Deque;
import java.util.LinkedList;
import java.util.List;

import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.finesys.common.dataPermission.factory.DataPermissionRuleFactory;
import org.finesys.common.dataPermission.local.DataPermissionRuleContextHolder;
import org.finesys.common.dataPermission.local.MappedStatementCache;
import org.finesys.common.dataPermission.rule.DataPermissionRule;
import org.finesys.common.mp.core.utils.MyBatisUtils;

import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;

import cn.hutool.core.collection.CollUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.BinaryExpression;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Function;
import net.sf.jsqlparser.expression.NotExpression;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.ExistsExpression;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.FromItem;
import net.sf.jsqlparser.statement.select.Join;
import net.sf.jsqlparser.statement.select.LateralSubSelect;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import net.sf.jsqlparser.statement.select.SelectExpressionItem;
import net.sf.jsqlparser.statement.select.SelectItem;
import net.sf.jsqlparser.statement.select.SetOperationList;
import net.sf.jsqlparser.statement.select.SubJoin;
import net.sf.jsqlparser.statement.select.SubSelect;
import net.sf.jsqlparser.statement.select.WithItem;


/**
 * 数据权限拦截器 通过数据权限规则 重写SQL的方式来实现
 * <p/>
 * 主要的SQL重写方法可见 {@link #builderExpression(Expression, Table)} 方法 参考
 * {@link com.baomidou.mybatisplus.extension.plugins.inner.TenantLineInnerInterceptor} 实现
 */
@RequiredArgsConstructor
@Slf4j
public class DataPermissionDatabaseInterceptor extends JsqlParserSupport implements InnerInterceptor {


    /**
     * 数据权限规则工厂接口 管理容器中配置的数据权限规则
     */
    private final DataPermissionRuleFactory dataPermissionRuleFactory;
    private final MappedStatementCache mappedStatementCache = new MappedStatementCache();

    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        //获取生效的数据规则
        List<DataPermissionRule> rules = dataPermissionRuleFactory.getDataPermissionRule();
        //如果无需重写，则跳过
        if (mappedStatementCache.noRewritable(ms, rules)) {
            return;
        }
        PluginUtils.MPBoundSql mpBoundSql = PluginUtils.mpBoundSql(boundSql);
        try {
            //初始化上下文
            DataPermissionRuleContextHolder.init(rules);
            //处理sql
            mpBoundSql.sql(parserSingle(mpBoundSql.sql(), null));
        } finally {
            addMappedStatementCache(ms);
            DataPermissionRuleContextHolder.clear();
        }
    }

    /**
     * 只处理 UPDATE / DELETE 场景 不处理 INSERT 场景
     */
    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        PluginUtils.MPStatementHandler mpStatementHandler = PluginUtils.mpStatementHandler(sh);
        MappedStatement mappedStatement = mpStatementHandler.mappedStatement();
        SqlCommandType sqlCommandType = mappedStatement.getSqlCommandType();
        if (sqlCommandType == SqlCommandType.UPDATE || sqlCommandType == SqlCommandType.DELETE) {
            //获取生效的数据权限规则
            List<DataPermissionRule> rules = dataPermissionRuleFactory.getDataPermissionRule();
            //如果无需重写则跳过
            if (mappedStatementCache.noRewritable(mappedStatement, rules)) {
                return;
            }

            PluginUtils.MPBoundSql mpBoundSql = mpStatementHandler.mPBoundSql();
            try {
                //初始化上下文
                DataPermissionRuleContextHolder.init(rules);
                //处理sql
                mpBoundSql.sql(parserMulti(mpBoundSql.sql(), null));
            } finally {
                addMappedStatementCache(mappedStatement);
                DataPermissionRuleContextHolder.clear();
            }
        }
    }


    @Override
    protected void processSelect(Select select, int index, String sql, Object obj) {
        processSelectBody(select.getSelectBody());
        List<WithItem> withItems = select.getWithItemsList();
        if (!CollUtil.isEmpty(withItems)) {
            withItems.forEach(this::processSelectBody);
        }
    }

    private void processSelectBody(SelectBody selectBody) {
        if (selectBody == null) {
            return;
        }
        if (selectBody instanceof PlainSelect) {
            processPlainSelect((PlainSelect) selectBody);
            return;
        }
        if (selectBody instanceof WithItem) {
            WithItem withItem = (WithItem) selectBody;
            processSelectBody(withItem.getSubSelect().getSelectBody());
            return;
        }
        SetOperationList operationList = (SetOperationList) selectBody;
        List<SelectBody> selectBodyList = operationList.getSelects();
        if (CollUtil.isNotEmpty(selectBodyList)) {
            selectBodyList.forEach(this::processSelectBody);
        }

    }

    private void processPlainSelect(PlainSelect plainSelect) {
        FromItem fromItem = plainSelect.getFromItem();
        Expression where = plainSelect.getWhere();
        processWhereSubSelect(where);
        if (fromItem instanceof Table) {
            Table fromTable = (Table) fromItem;
            plainSelect.setWhere(builderExpression(where, fromTable));
        } else {
            processFromItem(fromItem);
        }
        //
        List<SelectItem> selectItemList = plainSelect.getSelectItems();
        if (CollUtil.isNotEmpty(selectItemList)) {
            selectItemList.forEach(this::processSelectItem);
        }
        List<Join> joins = plainSelect.getJoins();
        if (CollUtil.isNotEmpty(joins)) {
            joins.forEach(this::processJoin);
        }

    }


    /**
     * 处理where条件内的子查询
     * <p>
     * 支持如下: 1. in 2. = 3. > 4. < 5. >= 6. <= 7. <> 8. EXISTS 9. NOT EXISTS
     * <p>
     * 前提条件: 1. 子查询必须放在小括号中 2. 子查询一般放在比较操作符的右边
     *
     * @param where where 条件
     */
    private void processWhereSubSelect(Expression where) {
        if (where == null) {
            return;
        }
        if (where instanceof FromItem) {
            processFromItem((FromItem) where);
            return;
        }
        if (where.toString().indexOf("SELECT") > 0) {
            //有子查询
            if (where instanceof BinaryExpression) {
                //比较符号and or
                BinaryExpression expression = (BinaryExpression) where;
                processWhereSubSelect(expression.getLeftExpression());
                processWhereSubSelect(expression.getRightExpression());
                return;
            }
            //in
            if (where instanceof InExpression) {
                InExpression expression = (InExpression) where;
                ItemsList itemsList = expression.getRightItemsList();
                if (itemsList instanceof SubSelect) {
                    processSelectBody(((SubSelect) itemsList).getSelectBody());
                }
                return;
            }
            //exists
            if (where instanceof ExistsExpression) {
                ExistsExpression expression = (ExistsExpression) where;
                processWhereSubSelect(expression.getRightExpression());
                return;
            }
            //not
            if (where instanceof NotExpression) {
                NotExpression expression = (NotExpression) where;
                processWhereSubSelect(expression.getExpression());
                return;
            }
            //
            if (where instanceof Parenthesis) {
                Parenthesis expression = (Parenthesis) where;
                processWhereSubSelect(expression.getExpression());
            }
        }

    }

    private void processSelectItem(SelectItem selectItem) {

        if (selectItem instanceof SelectExpressionItem) {
            SelectExpressionItem selectExpressionItem = (SelectExpressionItem) selectItem;
            if (selectExpressionItem.getExpression() instanceof SubSelect) {
                processSelectBody(((SubSelect) selectExpressionItem.getExpression()).getSelectBody());
                return;
            }
            if (selectExpressionItem.getExpression() instanceof Function) {
                processFunction((Function) selectExpressionItem.getExpression());
            }
        }
    }

    /**
     * 处理函数
     * <p>
     * 支持: 1. select fun(args..) 2. select fun1(fun2(args..),args..)
     * <p>
     * <p>
     * fixed gitee pulls/141
     * </p>
     *
     * @param function 函数
     */
    private void processFunction(Function function) {
        ExpressionList parameters = function.getParameters();
        if (parameters != null) {
            parameters.getExpressions().forEach(expression -> {
                if (expression instanceof SubSelect) {
                    processSelectBody(((SubSelect) expression).getSelectBody());
                    return;
                }
                if (expression instanceof Function) {
                    processFunction((Function) expression);
                }
            });
        }
    }

    /**
     * 处理子查询等
     */
    private void processFromItem(FromItem fromItem) {
        if (fromItem instanceof SubJoin) {
            SubJoin subJoin = (SubJoin) fromItem;
            if (subJoin.getJoinList() != null) {
                processJoins(subJoin.getJoinList());
            }
            if (subJoin.getLeft() != null) {
                processFromItem(subJoin.getLeft());
            }
            return;
        }

        if (fromItem instanceof SubSelect) {
            SubSelect subSelect = (SubSelect) fromItem;
            if (subSelect != null) {
                processSelectBody(subSelect.getSelectBody());
            }
            return;
        }

        if (fromItem instanceof LateralSubSelect) {
            LateralSubSelect lateralSubSelect = (LateralSubSelect) fromItem;
            if (lateralSubSelect.getSubSelect() != null) {
                SubSelect subSelect = lateralSubSelect.getSubSelect();
                if (subSelect.getSelectBody() != null) {
                    processSelectBody(subSelect.getSelectBody());
                }
            }
        }

    }

    private void processJoins(List<Join> joins) {
        // 对于 on 表达式写在最后的 join，需要记录下前面多个 on 的表名
        Deque<Table> tables = new LinkedList<>();
        for (Join join : joins) {
            //处理on表达式
            FromItem fromItem = join.getRightItem();
            if (fromItem instanceof Table) {
                Table fromTable = (Table) fromItem;
                //获取JOIN尾缀的on表达式列表
                Collection<Expression> originOnExpressions = join.getOnExpressions();
                //正常join on表达式只有一个，立刻处理
                if (originOnExpressions.size() == 1) {
                    processJoin(join);
                    continue;
                }
                tables.push(fromTable);
                //尾缀多个on表达式的时候统一处理
                if (originOnExpressions.size() > 1) {
                    Collection<Expression> onExpressions = new LinkedList<>();
                    for (Expression expression : originOnExpressions) {
                        Table currentTable = tables.poll();
                        onExpressions.add(builderExpression(expression, currentTable));
                    }
                    join.setOnExpressions(onExpressions);
                }
            } else {
                //处理右边连接的子表达式
                processFromItem(fromItem);
            }
        }

    }

    /**
     * 处理联接语句
     */
    private void processJoin(Join join) {
        if (join.getRightItem() instanceof Table) {
            Table fromTable = (Table) join.getRightItem();
            Expression originExpression = CollUtil.getFirst(join.getOnExpressions());
            originExpression = builderExpression(originExpression, fromTable);
            join.setOnExpressions(CollUtil.newArrayList(originExpression));
        }
    }


    private Expression builderExpression(Expression currentExpression, Table table) {

        // 获得 Table 对应的数据权限条件
        Expression equalsTo = buildDataPermissionExpression(table);
        // 如果没条件，则返回 currentExpression 默认
        if (equalsTo == null) {
            return currentExpression;
        }
        //表达式为空，则直接返回equalsTo
        if (currentExpression == null) {
            return equalsTo;
        }
        // 如果表达式为 Or，则需要 (currentExpression) AND equalsTo
        if (currentExpression instanceof OrExpression) {
            return new AndExpression(new Parenthesis(currentExpression), equalsTo);
        }
        // 如果表达式为 And，则直接返回 currentExpression AND equalsTo
        return new AndExpression(currentExpression, equalsTo);
    }

    private Expression buildDataPermissionExpression(Table table) {
        //生成条件
        Expression expressionAll = null;
        for (DataPermissionRule dataPermissionRule : DataPermissionRuleContextHolder.getRules()) {
            // 判断表名是否匹配
            if (!dataPermissionRule.getTableNames().contains(table.getName())) {
                continue;
            }
            // 如果有匹配的规则，说明可重写。
            // 为什么不是有 allExpression 非空才重写呢？在生成 column = value 过滤条件时，会因为 value 不存在，导致未重写。
            // 这样导致第一次无 value，被标记成无需重写；但是第二次有 value，此时会需要重写。
            DataPermissionRuleContextHolder.setRewrite(true);
            //单条件规则的条件
            String tableName = MyBatisUtils.getTableName(table);
            Expression oneExpression = dataPermissionRule.getExpression(tableName, table.getAlias());
            // 拼接到 allExpression 中
            expressionAll = expressionAll == null ? oneExpression : new AndExpression(expressionAll, oneExpression);
        }
        return expressionAll;
    }


    /**
     * 判断 SQL 是否重写 如果没有重写 则添加到 {@link MappedStatementCache} 中
     *
     * @param ms MappedStatement
     */
    private void addMappedStatementCache(MappedStatement ms) {
        if (DataPermissionRuleContextHolder.getRewrite()) {
            return;
        }
        //不重写
        mappedStatementCache.addNoRewritable(ms, DataPermissionRuleContextHolder.getRules());
    }
}
